{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Writing Custom Sample Parsers\n",
    "\n",
    "This recipe demonstrates how to write a [custom SampleParser](https://voxel51.com/docs/fiftyone/user_guide/sample_parsers.html#writing-a-custom-sampleparser) and use it to add samples in your custom format to a FiftyOne Dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "If you haven't already, install FiftyOne:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install fiftyone"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this receipe we'll use the [TorchVision Datasets](https://pytorch.org/vision/stable/datasets.html) library to download the [CIFAR-10 dataset](https://www.cs.toronto.edu/~kriz/cifar.html) to use as sample data to feed our custom parser.\n",
    "\n",
    "You can install the necessary packages, if necessary, as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install torch torchvision"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Writing a SampleParser\n",
    "\n",
    "FiftyOne provides a [SampleParser](https://voxel51.com/docs/fiftyone/api/fiftyone.utils.data.html#fiftyone.utils.data.parsers.SampleParser) interface that defines how it parses provided samples when methods such as [Dataset.add_labeled_images()](https://voxel51.com/docs/fiftyone/api/fiftyone.core.html#fiftyone.core.dataset.Dataset.add_labeled_images) and [Dataset.ingest_labeled_images()](https://voxel51.com/docs/fiftyone/api/fiftyone.core.html#fiftyone.core.dataset.Dataset.ingest_labeled_images) are used.\n",
    "\n",
    "`SampleParser` itself is an abstract interface; the concrete interface that you should implement is determined by the type of samples that you are importing. See [writing a custom SampleParser](https://voxel51.com/docs/fiftyone/user_guide/sample_parsers.html#writing-a-custom-sampleparser) for full details.\n",
    "\n",
    "In this recipe, we'll write a custom [LabeledImageSampleParser](https://voxel51.com/docs/fiftyone/api/fiftyone.utils.data.html#fiftyone.utils.data.parsers.LabeledImageSampleParser) that can parse labeled images from a [PyTorch Dataset](https://pytorch.org/docs/stable/data.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's the complete definition of the `SampleParser`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import fiftyone as fo\n",
    "import fiftyone.utils.data as foud\n",
    "\n",
    "\n",
    "class PyTorchClassificationDatasetSampleParser(foud.LabeledImageSampleParser):\n",
    "    \"\"\"Parser for image classification samples loaded from a PyTorch dataset.\n",
    "\n",
    "    This parser can parse samples from a ``torch.utils.data.DataLoader`` that\n",
    "    emits ``(img_tensor, target)`` tuples, where::\n",
    "\n",
    "        - `img_tensor`: is a PyTorch Tensor containing the image\n",
    "        - `target`: the integer index of the target class\n",
    "\n",
    "    Args:\n",
    "        classes: the list of class label strings\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, classes):\n",
    "        self.classes = classes\n",
    "\n",
    "    @property\n",
    "    def has_image_path(self):\n",
    "        \"\"\"Whether this parser produces paths to images on disk for samples\n",
    "        that it parses.\n",
    "        \"\"\"\n",
    "        return False\n",
    "\n",
    "    @property\n",
    "    def has_image_metadata(self):\n",
    "        \"\"\"Whether this parser produces\n",
    "        :class:`fiftyone.core.metadata.ImageMetadata` instances for samples\n",
    "        that it parses.\n",
    "        \"\"\"\n",
    "        return False\n",
    "\n",
    "    @property\n",
    "    def label_cls(self):\n",
    "        \"\"\"The :class:`fiftyone.core.labels.Label` class(es) returned by this\n",
    "        parser.\n",
    "\n",
    "        This can be any of the following:\n",
    "\n",
    "        -   a :class:`fiftyone.core.labels.Label` class. In this case, the\n",
    "            parser is guaranteed to return labels of this type\n",
    "        -   a list or tuple of :class:`fiftyone.core.labels.Label` classes. In\n",
    "            this case, the parser can produce a single label field of any of\n",
    "            these types\n",
    "        -   a dict mapping keys to :class:`fiftyone.core.labels.Label` classes.\n",
    "            In this case, the parser will return label dictionaries with keys\n",
    "            and value-types specified by this dictionary. Not all keys need be\n",
    "            present in the imported labels\n",
    "        -   ``None``. In this case, the parser makes no guarantees about the\n",
    "            labels that it may return\n",
    "        \"\"\"\n",
    "        return fo.Classification\n",
    "\n",
    "    def get_image(self):\n",
    "        \"\"\"Returns the image from the current sample.\n",
    "\n",
    "        Returns:\n",
    "            a numpy image\n",
    "        \"\"\"\n",
    "        img_tensor = self.current_sample[0]\n",
    "        return img_tensor.cpu().numpy()\n",
    "\n",
    "    def get_label(self):\n",
    "        \"\"\"Returns the label for the current sample.\n",
    "\n",
    "        Returns:\n",
    "            a :class:`fiftyone.core.labels.Label` instance, or a dictionary\n",
    "            mapping field names to :class:`fiftyone.core.labels.Label`\n",
    "            instances, or ``None`` if the sample is unlabeled\n",
    "        \"\"\"\n",
    "        target = self.current_sample[1]\n",
    "        return fo.Classification(label=self.classes[int(target)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that `PyTorchClassificationDatasetSampleParser` specifies `has_image_path == False` and `has_image_metadata == False`, because the PyTorch dataset directly provides the in-memory image, not its path on disk."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ingesting samples into a dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to use `PyTorchClassificationDatasetSampleParser`, we need a PyTorch Dataset from which to feed it samples.\n",
    "\n",
    "Let's use the [CIFAR-10 dataset](https://www.cs.toronto.edu/~kriz/cifar.html) from the [TorchVision Datasets](https://pytorch.org/docs/stable/torchvision/datasets.html) library:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz to /tmp/fiftyone/custom-parser/pytorch/cifar-10-python.tar.gz\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c5352424e69f494bb0d0ce851d254205",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=1.0, bar_style='info', max=1.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting /tmp/fiftyone/custom-parser/pytorch/cifar-10-python.tar.gz to /tmp/fiftyone/custom-parser/pytorch\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "\n",
    "\n",
    "# Downloads the test split of the CIFAR-10 dataset and prepares it for loading\n",
    "# in a DataLoader\n",
    "dataset = torchvision.datasets.CIFAR10(\n",
    "    \"/tmp/fiftyone/custom-parser/pytorch\",\n",
    "    train=False,\n",
    "    download=True,\n",
    "    transform=torchvision.transforms.ToTensor(),\n",
    ")\n",
    "classes = dataset.classes\n",
    "data_loader = torch.utils.data.DataLoader(dataset, batch_size=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can load the samples into the dataset. Since our custom sample parser declares `has_image_path == False`, we must use the [Dataset.ingest_labeled_images()](https://voxel51.com/docs/fiftyone/api/fiftyone.core.html#fiftyone.core.dataset.Dataset.ingest_labeled_images) method to load the samples into a FiftyOne dataset, which will write the individual images to disk as they are ingested so that FiftyOne can access them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 100% |███| 10000/10000 [6.7s elapsed, 0s remaining, 1.5K samples/s]      \n",
      "Loaded 10000 samples\n"
     ]
    }
   ],
   "source": [
    "dataset = fo.Dataset(\"cifar10-samples\")\n",
    "\n",
    "sample_parser = PyTorchClassificationDatasetSampleParser(classes)\n",
    "\n",
    "# The directory to use to store the individual images on disk\n",
    "dataset_dir = \"/tmp/fiftyone/custom-parser/fiftyone\"\n",
    "\n",
    "# Ingest the samples from the data loader\n",
    "dataset.ingest_labeled_images(data_loader, sample_parser, dataset_dir=dataset_dir)\n",
    "\n",
    "print(\"Loaded %d samples\" % len(dataset))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's inspect the contents of the dataset to verify that the samples were loaded as expected:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Name:           cifar10-samples\n",
      "Persistent:     False\n",
      "Num samples:    10000\n",
      "Tags:           []\n",
      "Sample fields:\n",
      "    filepath:     fiftyone.core.fields.StringField\n",
      "    tags:         fiftyone.core.fields.ListField(fiftyone.core.fields.StringField)\n",
      "    metadata:     fiftyone.core.fields.EmbeddedDocumentField(fiftyone.core.metadata.Metadata)\n",
      "    ground_truth: fiftyone.core.fields.StringField\n"
     ]
    }
   ],
   "source": [
    "# Print summary information about the dataset\n",
    "print(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<Sample: {\n",
      "    'dataset_name': 'cifar10-samples',\n",
      "    'id': '5f15aeab6d4e59654468a14e',\n",
      "    'filepath': '/tmp/fiftyone/custom-parser/fiftyone/000001.jpg',\n",
      "    'tags': BaseList([]),\n",
      "    'metadata': None,\n",
      "    'ground_truth': 'cat',\n",
      "}>\n",
      "<Sample: {\n",
      "    'dataset_name': 'cifar10-samples',\n",
      "    'id': '5f15aeab6d4e59654468a14f',\n",
      "    'filepath': '/tmp/fiftyone/custom-parser/fiftyone/000002.jpg',\n",
      "    'tags': BaseList([]),\n",
      "    'metadata': None,\n",
      "    'ground_truth': 'ship',\n",
      "}>\n",
      "<Sample: {\n",
      "    'dataset_name': 'cifar10-samples',\n",
      "    'id': '5f15aeab6d4e59654468a150',\n",
      "    'filepath': '/tmp/fiftyone/custom-parser/fiftyone/000003.jpg',\n",
      "    'tags': BaseList([]),\n",
      "    'metadata': None,\n",
      "    'ground_truth': 'ship',\n",
      "}>\n"
     ]
    }
   ],
   "source": [
    "# Print a few samples from the dataset\n",
    "print(dataset.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also verify that the ingested images were written to disk as expected:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total 0\r\n",
      "drwxr-xr-x  10002 voxel51  wheel   313K Jul 20 10:34 .\r\n",
      "drwxr-xr-x      4 voxel51  wheel   128B Jul 20 10:34 ..\r\n",
      "-rw-r--r--      1 voxel51  wheel     0B Jul 20 10:34 000001.jpg\r\n",
      "-rw-r--r--      1 voxel51  wheel     0B Jul 20 10:34 000002.jpg\r\n",
      "-rw-r--r--      1 voxel51  wheel     0B Jul 20 10:34 000003.jpg\r\n",
      "-rw-r--r--      1 voxel51  wheel     0B Jul 20 10:34 000004.jpg\r\n",
      "-rw-r--r--      1 voxel51  wheel     0B Jul 20 10:34 000005.jpg\r\n",
      "-rw-r--r--      1 voxel51  wheel     0B Jul 20 10:34 000006.jpg\r\n",
      "-rw-r--r--      1 voxel51  wheel     0B Jul 20 10:34 000007.jpg\r\n"
     ]
    }
   ],
   "source": [
    "!ls -lah /tmp/fiftyone/custom-parser/fiftyone | head -n 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adding samples to a dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If our `LabeledImageSampleParser` declared `has_image_path == True`, then we could use [Dataset.add_labeled_images()](https://voxel51.com/docs/fiftyone/api/fiftyone.core.html#fiftyone.core.dataset.Dataset.add_labeled_images) to add samples to FiftyOne datasets without creating a copy of the source images on disk.\n",
    "\n",
    "However, our sample parser does not provide image paths, so an informative error message is raised if we try to use it in an unsupported way:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "Sample parser must have `has_image_path == True` to add its samples to the dataset",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-6-a3d739e371af>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0;31m# Won't work because our SampleParser does not provide paths to its source images on disk\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 6\u001b[0;31m \u001b[0mdataset\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0madd_labeled_images\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata_loader\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msample_parser\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m~/dev/fiftyone/fiftyone/core/dataset.py\u001b[0m in \u001b[0;36madd_labeled_images\u001b[0;34m(self, samples, sample_parser, label_field, tags, expand_schema)\u001b[0m\n\u001b[1;32m    729\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0msample_parser\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mhas_image_path\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    730\u001b[0m             raise ValueError(\n\u001b[0;32m--> 731\u001b[0;31m                 \u001b[0;34m\"Sample parser must have `has_image_path == True` to add its \"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    732\u001b[0m                 \u001b[0;34m\"samples to the dataset\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    733\u001b[0m             )\n",
      "\u001b[0;31mValueError\u001b[0m: Sample parser must have `has_image_path == True` to add its samples to the dataset"
     ]
    }
   ],
   "source": [
    "dataset = fo.Dataset()\n",
    "\n",
    "sample_parser = PyTorchClassificationDatasetSampleParser(classes)\n",
    "\n",
    "# Won't work because our SampleParser does not provide paths to its source images on disk\n",
    "dataset.add_labeled_images(data_loader, sample_parser)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cleanup\n",
    "\n",
    "You can cleanup the files generated by this recipe by running:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm -rf /tmp/fiftyone"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
